package dan.frontend
import chisel3._
import chisel3.util._
import chisel3.util.random.LFSR
import dan.common.Consts._
import dan.common.SDPRam
import dan.utility._
import org.fusesource.jansi.internal.Kernel32.INPUT_RECORD
import dan.utility.ImplicitCast.uintToBitPat

class BTBEntry extends Bundle with HasBTBParam{
    val target = UInt(targetBits.W)
}

class BTBInfo extends Bundle with HasBTBParam{
    val tag = UInt(tagBits.W)
    val isBr = Bool()
    // 与MicroBTB不同，2位计数器留到BIM模块中
}

class BTBMeta extends Bundle with HasBTBParam{
    val hit = Bool()
    val wayIdx = UInt(log2Ceil(wayNum).W)
}

class BTBIO extends Bundle with HasBTBParam{
    val s0Valid = Input(Bool())
    val s0PC = Input(UInt(vaBits.W))
    val s0HashedPC = Input(UInt(vaBits.W))
    val s2Target = Output(Vec(fetchWidth, Valid(UInt(targetBits.W))))
    val s2Taken = Output(Vec(fetchWidth, Bool()))
    val s2IsBr = Output(Vec(fetchWidth, Bool()))
    val s2IsJal = Output(Vec(fetchWidth, Bool()))
    val s3Meta = Output(Vec(fetchWidth, new BTBMeta()))
    val s1Update = Input(Valid(new PredictionUpdate()))
}

class BTB extends  Module with HasBTBParam{
    val io = IO(new BTBIO())
    // TODO SDPRam
    val info = Seq.fill(wayNum){Module(new SDPRam(setNum, new BTBInfo(), fetchWidth))}
    val targets = Seq.fill(wayNum){Module(new SDPRam(setNum, new BTBEntry(), fetchWidth))}
    
    val resetEnable = RegInit(false.B)
    val resetIdx = RegInit(0.U(log2Ceil(setNum).W))
    // 需要手动reset
    when(resetEnable){
        resetIdx := resetIdx + 1.U
    }
    when(resetIdx === (setNum - 1).U){
        resetEnable := false.B
    }
    val s0SetIdx = getSetIdx(io.s0HashedPC)
    val s0Tag = getTag(io.s0HashedPC)
    val s1Valid = RegNext(io.s0Valid)
    val s1PC = RegNext(io.s0PC)
    val s1Tag = RegNext(s0Tag)
    val s1Targets = VecInit(targets.map(
        t => {
            t.io.raddr := s0SetIdx
            t.io.rdata
        }
    ))
    val s1Info = VecInit(info.map(
        i => {
            i.io.raddr := s0SetIdx
            i.io.rdata
        }
    ))
    val s1HitVec = VecInit((0 until fetchWidth) map{
        i => VecInit((0 until wayNum) map {
            // tag相等
            j => IsEqual(s1Info(j)(i).tag, s1Tag.asUInt)
        })   
    })
    val s1GetHit = s1HitVec.map(_.asUInt.orR)
    // why use PriorityEncoder instead of OHToUInt?
    val s1HitWay = s1HitVec.map(vec => PriorityEncoder(vec))
    // 读出预测结果
    // TODO
    for(i <- 0 until fetchWidth){
        val s1GiveResp = !resetEnable && s1Valid && s1GetHit(i)
        val s1EntryInfo = s1Info(s1HitWay(i))(i)
        val s1EntryTarget = s1Targets(s1HitWay(i))(i)

        val s2IsBr = RegNext(s1GiveResp && s1EntryInfo.isBr)
        io.s2IsBr(i) := s2IsBr
        val s2IsJal = RegNext(s1GiveResp && !s1EntryInfo.isBr)
        io.s2IsJal(i) := s2IsJal
        io.s2Taken(i) := s2IsJal
        
        io.s2Target(i).valid := RegNext(s1GiveResp)
        io.s2Target(i).bits := RegNext(s1EntryTarget.target) // ?
    }
    // 随机替换
    val replaceEnable = s1Valid && !s1GetHit.reduce(_||_)
    val replaceWay = LFSR(wayNum, replaceEnable)(log2Ceil(wayNum)-1, 0)
    val s1Meta = Wire(Vec(fetchWidth, new BTBMeta()))
    for(i <- 0 until fetchWidth){
        s1Meta(i).hit := s1GetHit(i)
        s1Meta(i).wayIdx := Mux(s1GetHit(i), s1HitWay(i), replaceWay)
    }
    io.s3Meta := RegNext(RegNext(s1Meta))

    // 更新
    val s1Update = io.s1Update
    val s1UpdateHashedPC = mixHILO(s1Update.bits.pc)
    val s1UpdateTag = getTag(s1UpdateHashedPC)
    val s1UpdateSetIdx = getSetIdx(s1UpdateHashedPC)
    val s1UpdateCFIIdx = s1Update.bits.idx.bits
    val s1UpdateMeta = VecInit(s1Update.bits.meta.map(_.btbMeta))
    val s1UpdateWayIdx = VecInit(s1UpdateMeta.map(_.wayIdx))
    val s1UpdateFinalWay = s1UpdateWayIdx(s1UpdateCFIIdx)
    // 更新的一行（某个way的某个set），包含一个取值pack的4条指令，用相同的tag
    val s1UpdateInfo = Wire(Vec(fetchWidth, new BTBInfo()))
    for(i <- 0 until fetchWidth){
        s1UpdateInfo(i).tag := s1UpdateTag
        s1UpdateInfo(i).isBr := s1Update.bits.brMask(i)
    }
    val s1UpdateTarget = Wire(new BTBEntry())
    s1UpdateTarget.target := getTarget(s1Update.bits.target)
    // 需要更新的位置（CFI idx） & 是否需要更新（全0/1的使能向量）
    val s1TargetsWrMask = UIntToOH(s1UpdateCFIIdx) & Fill(fetchWidth, s1Update.valid && s1Update.bits.idx.valid && s1Update.bits.taken)
    val s1InfoWrMask = s1TargetsWrMask
    for(i <- 0 until wayNum){
        val updatable = s1UpdateFinalWay === i.U
        info(i).io.wen := resetEnable || updatable
        info(i).io.waddr := Mux(resetEnable, resetIdx, s1UpdateSetIdx)
        info(i).io.wdata := Mux(resetEnable, VecInit(Seq.fill(fetchWidth){0.U.asTypeOf(new BTBInfo())}), s1UpdateInfo)
        info(i).io.wstrobe := Mux(resetEnable, ~0.U(fetchWidth.W), s1InfoWrMask.asUInt)
        targets(i).io.wen := resetEnable || updatable
        targets(i).io.waddr := Mux(resetEnable, resetIdx, s1UpdateSetIdx)
        targets(i).io.wdata := VecInit(Seq.fill(fetchWidth) {Mux(resetEnable, 0.U.asTypeOf(new BTBEntry()), s1UpdateTarget)})
        targets(i).io.wstrobe := Mux(resetEnable, ~0.U(fetchWidth.W), s1TargetsWrMask.asUInt)
    }
}


